import torch
from pytorch_lightning import LightningModule
from torch.nn import Module, Linear, ReLU, Sequential, MSELoss


class WzzModel(Module):
    def __init__(self):
        super().__init__()
        self.encode = Sequential(
            Linear(in_features=29, out_features=12),
            ReLU(),
            Linear(in_features=12, out_features=4),
            ReLU()
        )
        self.decode = Sequential(
            Linear(in_features=4, out_features=12),
            ReLU(),
            Linear(in_features=12, out_features=29)
        )

    def forward(self, x):
        encode = self.encode(x)
        decode = self.decode(encode)
        return decode


class WzzModel_LM(LightningModule):
    def __init__(self):
        super().__init__()
        self.net = WzzModel()
        self.loss_fn = MSELoss()

    def forward(self, x):
        return self.net(x)

    def training_step(self, batch, batch_idx):
        x = batch[0]
        pred = self(x)
        loss = self.loss_fn(pred, x)
        self.log('loss', loss)
        return loss

    def configure_optimizers(self):
        return torch.optim.Adam(params=self.parameters(), lr=0.001)
